from fastmcp import FastMCP
import asyncio
from fastmcp import Client
from typing import Annotated
from datetime import datetime
from dateutil.relativedelta import relativedelta
from openai import OpenAI

import os
os.environ["OPENAI_API_KEY"]=""

config = {
    "backend_url": "",
    "quick_think_llm": "gpt-5"  #o3-deep-research, o4-mini-deep-research, or gpt-5.
}

mcp = FastMCP('demo.mcp')

@mcp.tool()
def get_stock_news(
    query: Annotated[str, "Search query, typically the company name or ticker (e.g., 'Apple' or 'AAPL')"],
    start_date: Annotated[str, "Start date in yyyy-mm-dd format (e.g., '2025-10-01')"],
    end_date: Annotated[str, "End date in yyyy-mm-dd format (e.g., '2025-10-08')"]
) -> str:
    """
    Retrieves social media posts related to a stock or company within a specified date range.

    Input Parameters:
        query: The search query, typically the company name or ticker symbol (e.g., 'AAPL' for Apple).
        start_date: The start date for the search in yyyy-mm-dd format. Must be a valid date.
        end_date: The end date for the search in yyyy-mm-dd format. Must be a valid date and after start_date.

    Output:
        result: A string containing the text of social media posts found for the query within the specified period.
                Returns an error message if the date format is invalid or no data is found.
    """
    try:
        datetime.strptime(start_date, "%Y-%m-%d")
        datetime.strptime(end_date, "%Y-%m-%d")
    except ValueError:
        return f"Invalid date format for start_date '{start_date}' or end_date '{end_date}'. Use yyyy-mm-dd."

    client = OpenAI(base_url=config["backend_url"])

    response = client.responses.create(
        model=config["quick_think_llm"],
        input=[
            {
                "role": "system",
                "content": [
                    {
                        "type": "input_text",
                        "text": f"Can you search Social Media for {query} from {start_date} to {end_date}? Make sure you only get the data posted during that period."
                    }
                ]
            }
        ],
        text={"format": {"type": "text"}},
        reasoning={},
        tools=[
            {
                "type": "web_search_preview",
                "user_location": {"type": "approximate"},
                "search_context_size": "low"
            }
        ],
        temperature=1,
        max_output_tokens=4096,
        top_p=1,
        store=True
    )

    return response.output[1].content[0].text

@mcp.tool()
def get_global_news(
    curr_date: Annotated[str, "Current date in yyyy-mm-dd format (e.g., '2025-10-09')"],
    look_back_days: Annotated[int, "Number of days to look back for news (default: 7)"] = 7,
    limit: Annotated[int, "Maximum number of articles to return (default: 5)"] = 5
) -> str:
    """
    Retrieves global or macroeconomic news relevant for trading within a specified period.

    Input Parameters:
        curr_date: The current date in yyyy-mm-dd format, marking the end of the search period.
        look_back_days: Number of days to look back from curr_date for news articles. Defaults to 7.
        limit: Maximum number of articles to return. Defaults to 5.

    Output:
        result: A string containing the text of global or macroeconomic news articles relevant for trading.
                Returns an error message if the date format is invalid or no data is found.
    """
    try:
        datetime.strptime(curr_date, "%Y-%m-%d")
    except ValueError:
        return f"Invalid date format for curr_date '{curr_date}'. Use yyyy-mm-dd."

    client = OpenAI(base_url=config["backend_url"])

    response = client.responses.create(
        model=config["quick_think_llm"],
        input=[
            {
                "role": "system",
                "content": [
                    {
                        "type": "input_text",
                        "text": f"Can you search global or macroeconomics news from {look_back_days} days before {curr_date} to {curr_date} that would be informative for trading purposes? Make sure you only get the data posted during that period. Limit the results to {limit} articles."
                    }
                ]
            }
        ],
        text={"format": {"type": "text"}},
        reasoning={},
        tools=[
            {
                "type": "web_search_preview", # These tasks typically use models like o3-deep-research, o4-mini-deep-research, or gpt-5 with reasoning level set to high
                "user_location": {"type": "approximate"},
                "search_context_size": "low"
            }
        ],
        temperature=1,
        max_output_tokens=4096,
        top_p=1,
        store=True
    )

    return response.output[1].content[0].text

@mcp.tool()
def get_fundamentals(
    ticker: Annotated[str, "Stock ticker symbol (e.g., 'AAPL' for Apple)"],
    curr_date: Annotated[str, "Current date in yyyy-mm-dd format (e.g., '2025-10-09')"]
) -> str:
    """
    Retrieves fundamental data discussions for a stock within the month before and including the current date.

    Input Parameters:
        ticker: The stock ticker symbol (e.g., 'AAPL' for Apple).
        curr_date: The current date in yyyy-mm-dd format, marking the end of the search period.

    Output:
        result: A string containing a table of fundamental data (e.g., PE, PS, Cash Flow) for the ticker.
                Returns an error message if the date format is invalid or no data is found.
    """
    try:
        datetime.strptime(curr_date, "%Y-%m-%d")
    except ValueError:
        return f"Invalid date format for curr_date '{curr_date}'. Use yyyy-mm-dd."

    client = OpenAI(base_url=config["backend_url"])

    response = client.responses.create(
        model=config["quick_think_llm"],
        input=[
            {
                "role": "system",
                "content": [
                    {
                        "type": "input_text",
                        "text": f"Can you search Fundamental for discussions on {ticker} during of the month before {curr_date} to the month of {curr_date}. Make sure you only get the data posted during that period. List as a table, with PE/PS/Cash flow/ etc"
                    }
                ]
            }
        ],
        text={"format": {"type": "text"}},
        reasoning={},
        tools=[
            {
                "type": "web_search_preview",
                "user_location": {"type": "approximate"},
                "search_context_size": "low"
            }
        ],
        temperature=1,
        max_output_tokens=4096,
        top_p=1,
        store=True
    )

    return response.output[1].content[0].text

OPENAI_STOCK_NEWS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_stock_news",
        "description": "Retrieves social media posts related to a stock or company within a specified date range.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search query, typically the company name or ticker symbol (e.g., 'AAPL' for Apple)."
                },
                "start_date": {
                    "type": "string",
                    "description": "The start date for the search in yyyy-mm-dd format. Must be a valid date."
                },
                "end_date": {
                    "type": "string",
                    "description": "The end date for the search in yyyy-mm-dd format. Must be a valid date and after start_date."
                }
            },
            "required": ["query", "start_date", "end_date"]
        }
    }
}

OPENAI_GLOBAL_NEWS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_global_news",
        "description": "Retrieves global or macroeconomic news relevant for trading within a specified period.",
        "parameters": {
            "type": "object",
            "properties": {
                "curr_date": {
                    "type": "string",
                    "description": "The current date in yyyy-mm-dd format, marking the end of the search period."
                },
                "look_back_days": {
                    "type": "integer",
                    "description": "Number of days to look back from curr_date for news articles.",
                    "default": 7
                },
                "limit": {
                    "type": "integer",
                    "description": "Maximum number of articles to return.",
                    "default": 5
                }
            },
            "required": ["curr_date"]
        }
    }
}

OPENAI_FUNDAMENTALS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_fundamentals",
        "description": "Retrieves fundamental data discussions for a stock within the month before and including the current date.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'AAPL' for Apple)."
                },
                "curr_date": {
                    "type": "string",
                    "description": "The current date in yyyy-mm-dd format, marking the end of the search period."
                }
            },
            "required": ["ticker", "curr_date"]
        }
    }
}

async def main():
    """
    """
    client = Client(mcp)
    async with client:
        # List available tools for debugging/verification
        tools = await client.list_tools()
        print('Available tools:', tools)

        # Call the stock news tool
        news_result = await client.call_tool('get_stock_news', {
            'query': 'AAPL',
            'start_date': '2025-10-01',
            'end_date': '2025-10-08'
        })
        print('Stock News Result:', news_result)

        # Call the global news tool
        global_news_result = await client.call_tool('get_global_news', {
            'curr_date': '2025-10-09',
            'look_back_days': 7,
            'limit': 5
        })
        print('Global News Result:', global_news_result)

        # Call the fundamentals tool
        fundamentals_result = await client.call_tool('get_fundamentals', {
            'ticker': 'AAPL',
            'curr_date': '2025-10-09'
        })
        print('Fundamentals Result:', fundamentals_result)

if __name__ == '__main__':
    asyncio.run(main())


